﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;

namespace CloudStorageLight.WpfClient
{
    /// <summary>
    /// A partial wrapper class for Restart Manager API
    /// </summary>
    /// <remarks>
    /// This class is based on http://msdn.microsoft.com/en-us/magazine/cc163450.aspx
    /// To use this class, OS has to be Windows Vista or newer.
    /// </remarks>
    public static class RestartManager
    {
        #region "Win32"

        // Start a Restart Manager session.
        [DllImport("Rstrtmgr.dll", CharSet = CharSet.Unicode)]
        private static extern uint RmStartSession(
            out uint pSessionHandle,
            int dwSessionFlags,
            string strSessionKey);

        // End a Restart Manager session.
        [DllImport("Rstrtmgr.dll")]
        private static extern uint RmEndSession(uint pSessionHandle);

        // Register target files to a Restart Manager session.
        [DllImport("Rstrtmgr.dll", CharSet = CharSet.Unicode)]
        private static extern uint RmRegisterResources(
            uint pSessionHandle,
            uint nFiles,
            string[] rgsFilenames,
            uint nApplications,
            RM_UNIQUE_PROCESS[] rgApplications,
            uint nServices,
            string[] rgsServiceNames);

        // Get processes using target files with a Restart Manager session.
        [DllImport("Rstrtmgr.dll")]
        private static extern uint RmGetList(
            uint pSessionHandle,
            out uint pnProcInfoNeeded,
            ref uint pnProcInfo,
            [In, Out] RM_PROCESS_INFO[] rgAffectedApps,
            ref uint lpdwRebootReasons);

        private const uint RmRebootReasonNone = 0U;

        private const int CCH_RM_MAX_APP_NAME = 255;
        private const int CCH_RM_MAX_SVC_NAME = 63;

        private const uint ERROR_SUCCESS = 0U;
        private const uint ERROR_MORE_DATA = 234U;

        [StructLayout(LayoutKind.Sequential)]
        private struct RM_UNIQUE_PROCESS
        {
            public uint dwProcessId;
            public System.Runtime.InteropServices.ComTypes.FILETIME ProcessStartTime;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        private struct RM_PROCESS_INFO
        {
            public RM_UNIQUE_PROCESS Process;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_APP_NAME + 1)]
            public string strAppName;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_SVC_NAME + 1)]
            public string strServiceShortName;
            public RM_APP_TYPE ApplicationType;
            public uint AppStatus;
            public uint TSSessionId;
            [MarshalAs(UnmanagedType.Bool)]
            public bool bRestartable;
        }

        private enum RM_APP_TYPE
        {
            /// <summary>
            /// The application cannot be classified as any other type.
            /// </summary>
            RmUnknownApp = 0,

            /// <summary>
            /// A Windows application run as a stand-alone process that displays a top-level window.
            /// </summary>
            RmMainWindow = 1,

            /// <summary>
            /// A Windows application that does not run as a stand-alone process and does not display a top-level window.
            /// </summary>
            RmOtherWindow = 2,

            /// <summary>
            /// The application is a Windows service.
            /// </summary>
            RmService = 3,

            /// <summary>
            /// The application is Windows Explorer.
            /// </summary>
            RmExplorer = 4,

            /// <summary>
            /// The application is a stand-alone console application.
            /// </summary>
            RmConsole = 5,

            /// <summary>
            /// The process may be a critical process and cannot be shut down.
            /// </summary>
            RmCritical = 1000
        }

        #endregion


        /// <summary>
        /// Check if any process is using a specified file.
        /// </summary>
        /// <param name="filePath">Path of target file</param>
        /// <returns>True if using</returns>
        public static bool IsProcessesUsingFile(string filePath)
        {
            bool isUsing = false;
            foreach (var proc in EnumerateProcessesUsingFiles(new[] { filePath }))
            {
                if (proc == null || proc.Id == Process.GetCurrentProcess().Id)
                    continue;

                isUsing = true;
                proc.Dispose();
            }

            return isUsing;
        }

        /// <summary>
        /// Enumerate processes using a specified file.
        /// </summary>
        /// <param name="filePath">Path of target file</param>
        /// <returns>Processes using target file</returns>
        /// <remarks>Caller is responsible for disposing the processes.</remarks>
        public static IEnumerable<Process> EnumerateProcessesUsingFile(string filePath)
        {
            return EnumerateProcessesUsingFiles(new[] { filePath });
        }

        /// <summary>
        /// Enumerate processes using specified files.
        /// </summary>
        /// <param name="filePaths">Paths of target files</param>
        /// <returns>Processes using target files</returns>
        /// <remarks>Caller is responsible for disposing the processes.</remarks>
        public static IEnumerable<Process> EnumerateProcessesUsingFiles(string[] filePaths)
        {
            if ((filePaths == null) || !filePaths.Any())
                yield break;

            if (!OsVersion.IsVistaOrNewer)
                yield break;

            uint sessionHandle = 0U; // Handle to Restart Manager session

            try
            {
                // Start a Restart Manager session.
                var result1 = RmStartSession(
                    out sessionHandle,
                    0,
                    Guid.NewGuid().ToString("N"));

                if (result1 != ERROR_SUCCESS)
                    throw new Win32Exception("Failed to start a Restart Manager session.");

                // Register target files to the session.
                var result2 = RmRegisterResources(
                    sessionHandle,
                    (uint)filePaths.Length,
                    filePaths,
                    0U,
                    null,
                    0U,
                    null);

                if (result2 != ERROR_SUCCESS)
                    throw new Win32Exception("Failed to register target files to a Restart Manager session.");

                // Get processes using target files with the session.
                uint pnProcInfoNeeded = 0U;
                uint pnProcInfo = 0U;
                RM_PROCESS_INFO[] rgAffectedApps = null;
                var lpdwRebootReasons = RmRebootReasonNone;

                uint result3;

                do
                {
                    result3 = RmGetList(
                        sessionHandle,
                        out pnProcInfoNeeded,
                        ref pnProcInfo,
                        rgAffectedApps,
                        ref lpdwRebootReasons);

                    switch (result3)
                    {
                        case ERROR_SUCCESS: // The size of RM_PROCESS_INFO array is appropriate.
                            if (pnProcInfo == 0)
                                break;

                            // Yield the processes.
                            foreach (var app in rgAffectedApps)
                            {
                                Process proc = null;
                                try
                                {
                                    proc = Process.GetProcessById((int)app.Process.dwProcessId);
                                }
                                catch (ArgumentException)
                                {
                                    // None (In case the process is no longer running).
                                }

                                if (proc != null)
                                    yield return proc;
                            }
                            break;

                        case ERROR_MORE_DATA: // The size of RM_PROCESS_INFO array is not enough.
                            // Set RM_PROCESS_INFO array to store the processes.
                            rgAffectedApps = new RM_PROCESS_INFO[(int)pnProcInfoNeeded];
                            pnProcInfo = (uint)rgAffectedApps.Length;
                            break;

                        default:
                            throw new Win32Exception("Failed to get processes using target files with a Restart Manager session.");
                    }
                }
                while (result3 != ERROR_SUCCESS);
            }
            finally
            {
                // End the session.
                RmEndSession(sessionHandle);
            }
        }
    }

    public static class OsVersion
    {
        private static readonly Version ver = Environment.OSVersion.Version;

        /// <summary>
        /// Whether OS is Windows Vista or newer
        /// </summary>
        /// <remarks>Windows Vista = version 6.0</remarks>
        public static bool IsVistaOrNewer
        {
            get { return (6 <= ver.Major); }
        }
    }
}
